"""
FastAPI后端接口
提供RESTful API服务
"""

from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
import os
import sys
import uuid
from pathlib import Path
import logging

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from src.models.text2video import TextToVideoGenerator
from src.models.style_transfer import StyleTransfer
from src.models.video_editing import VideoEditor
from src.utils.text_utils import TextProcessor

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 初始化模型
text_processor = TextProcessor()
generator = None
transfer = None
editor = None

# 创建FastAPI应用
app = FastAPI(
    title="AI动画生成系统API",
    description="基于深度学习的AI动画生成系统API接口",
    version="1.0.0"
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 数据模型
class TextToVideoRequest(BaseModel):
    prompt: str
    style: str = "realistic"
    duration: float = 3.0
    fps: int = 8
    resolution: str = "512x512"

class StyleTransferRequest(BaseModel):
    style: str
    strength: float = 0.8

class VideoEditRequest(BaseModel):
    operations: List[str]
    start_time: Optional[float] = None
    end_time: Optional[float] = None

class TaskResponse(BaseModel):
    task_id: str
    status: str
    message: str

# 任务存储
tasks = {}

def load_models():
    """加载模型"""
    global generator, transfer, editor
    
    try:
        if generator is None:
            generator = TextToVideoGenerator()
        if transfer is None:
            transfer = StyleTransfer()
        if editor is None:
            editor = VideoEditor()
        
        logger.info("模型加载成功")
        return True
    except Exception as e:
        logger.error(f"模型加载失败: {e}")
        return False

@app.on_event("startup")
async def startup_event():
    """应用启动事件"""
    logger.info("AI动画生成系统API启动中...")
    load_models()

@app.get("/")
async def root():
    """根路径"""
    return {
        "message": "AI动画生成系统API",
        "version": "1.0.0",
        "status": "running"
    }

@app.get("/health")
async def health_check():
    """健康检查"""
    models_loaded = all([generator, transfer, editor])
    return {
        "status": "healthy" if models_loaded else "unhealthy",
        "models_loaded": models_loaded,
        "generator": generator is not None,
        "transfer": transfer is not None,
        "editor": editor is not None
    }

@app.post("/api/text2video", response_model=TaskResponse)
async def text_to_video(request: TextToVideoRequest, background_tasks: BackgroundTasks):
    """文本到视频生成"""
    try:
        # 生成任务ID
        task_id = str(uuid.uuid4())
        
        # 初始化任务
        tasks[task_id] = {
            "status": "processing",
            "message": "正在生成视频...",
            "result": None
        }
        
        # 后台处理
        background_tasks.add_task(
            process_text_to_video,
            task_id,
            request.prompt,
            request.style,
            request.duration,
            request.fps,
            request.resolution
        )
        
        return TaskResponse(
            task_id=task_id,
            status="processing",
            message="任务已提交，正在处理中"
        )
        
    except Exception as e:
        logger.error(f"文本到视频生成失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

async def process_text_to_video(task_id: str, prompt: str, style: str, 
                              duration: float, fps: int, resolution: str):
    """处理文本到视频生成任务"""
    try:
        # 确保模型已加载
        if not load_models():
            tasks[task_id] = {
                "status": "failed",
                "message": "模型加载失败",
                "result": None
            }
            return
        
        # 优化提示词
        optimized_prompt = text_processor.optimize_prompt(prompt, style)
        
        # 生成视频
        output_path = generator.generate(
            prompt=optimized_prompt,
            style=style,
            duration=duration,
            fps=fps,
            resolution=tuple(map(int, resolution.split('x')))
        )
        
        # 更新任务状态
        tasks[task_id] = {
            "status": "completed",
            "message": "视频生成成功",
            "result": output_path
        }
        
        logger.info(f"任务 {task_id} 完成: {output_path}")
        
    except Exception as e:
        logger.error(f"任务 {task_id} 失败: {e}")
        tasks[task_id] = {
            "status": "failed",
            "message": f"生成失败: {str(e)}",
            "result": None
        }

@app.post("/api/style_transfer", response_model=TaskResponse)
async def style_transfer(
    file: UploadFile = File(...),
    request: StyleTransferRequest = None
):
    """风格迁移"""
    try:
        # 验证文件
        if not file.filename.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
            raise HTTPException(status_code=400, detail="不支持的文件格式")
        
        # 生成任务ID
        task_id = str(uuid.uuid4())
        
        # 保存上传文件
        input_path = f"data/temp/input_{task_id}_{file.filename}"
        os.makedirs("data/temp", exist_ok=True)
        
        with open(input_path, "wb") as f:
            content = await file.read()
            f.write(content)
        
        # 初始化任务
        tasks[task_id] = {
            "status": "processing",
            "message": "正在处理风格迁移...",
            "result": None,
            "input_path": input_path
        }
        
        # 后台处理
        from fastapi import BackgroundTasks
        background_tasks = BackgroundTasks()
        background_tasks.add_task(
            process_style_transfer,
            task_id,
            input_path,
            request.style if request else "cartoon",
            request.strength if request else 0.8
        )
        
        return TaskResponse(
            task_id=task_id,
            status="processing",
            message="任务已提交，正在处理中"
        )
        
    except Exception as e:
        logger.error(f"风格迁移失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

async def process_style_transfer(task_id: str, input_path: str, style: str, strength: float):
    """处理风格迁移任务"""
    try:
        # 确保模型已加载
        if not load_models():
            tasks[task_id] = {
                "status": "failed",
                "message": "模型加载失败",
                "result": None
            }
            return
        
        # 执行风格迁移
        output_path = transfer.transfer(
            input_path=input_path,
            style=style,
            strength=strength
        )
        
        # 更新任务状态
        tasks[task_id] = {
            "status": "completed",
            "message": "风格迁移成功",
            "result": output_path
        }
        
        logger.info(f"任务 {task_id} 完成: {output_path}")
        
    except Exception as e:
        logger.error(f"任务 {task_id} 失败: {e}")
        tasks[task_id] = {
            "status": "failed",
            "message": f"风格迁移失败: {str(e)}",
            "result": None
        }

@app.post("/api/video_edit", response_model=TaskResponse)
async def video_edit(
    file: UploadFile = File(...),
    request: VideoEditRequest = None
):
    """视频编辑"""
    try:
        # 验证文件
        if not file.filename.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
            raise HTTPException(status_code=400, detail="不支持的文件格式")
        
        # 生成任务ID
        task_id = str(uuid.uuid4())
        
        # 保存上传文件
        input_path = f"data/temp/input_{task_id}_{file.filename}"
        os.makedirs("data/temp", exist_ok=True)
        
        with open(input_path, "wb") as f:
            content = await file.read()
            f.write(content)
        
        # 初始化任务
        tasks[task_id] = {
            "status": "processing",
            "message": "正在编辑视频...",
            "result": None,
            "input_path": input_path
        }
        
        # 后台处理
        from fastapi import BackgroundTasks
        background_tasks = BackgroundTasks()
        background_tasks.add_task(
            process_video_edit,
            task_id,
            input_path,
            request.operations if request else ["trim"],
            request.start_time if request else None,
            request.end_time if request else None
        )
        
        return TaskResponse(
            task_id=task_id,
            status="processing",
            message="任务已提交，正在处理中"
        )
        
    except Exception as e:
        logger.error(f"视频编辑失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

async def process_video_edit(task_id: str, input_path: str, operations: List[str],
                           start_time: Optional[float], end_time: Optional[float]):
    """处理视频编辑任务"""
    try:
        # 确保模型已加载
        if not load_models():
            tasks[task_id] = {
                "status": "failed",
                "message": "模型加载失败",
                "result": None
            }
            return
        
        # 执行视频编辑
        output_path = editor.edit(
            input_path=input_path,
            operations=operations,
            start_time=start_time,
            end_time=end_time
        )
        
        # 更新任务状态
        tasks[task_id] = {
            "status": "completed",
            "message": "视频编辑成功",
            "result": output_path
        }
        
        logger.info(f"任务 {task_id} 完成: {output_path}")
        
    except Exception as e:
        logger.error(f"任务 {task_id} 失败: {e}")
        tasks[task_id] = {
            "status": "failed",
            "message": f"视频编辑失败: {str(e)}",
            "result": None
        }

@app.get("/api/status/{task_id}")
async def get_task_status(task_id: str):
    """获取任务状态"""
    if task_id not in tasks:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    task = tasks[task_id]
    return {
        "task_id": task_id,
        "status": task["status"],
        "message": task["message"],
        "result": task.get("result")
    }

@app.get("/api/download/{task_id}")
async def download_result(task_id: str):
    """下载结果文件"""
    if task_id not in tasks:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    task = tasks[task_id]
    if task["status"] != "completed" or not task.get("result"):
        raise HTTPException(status_code=400, detail="任务未完成或结果不存在")
    
    result_path = task["result"]
    if not os.path.exists(result_path):
        raise HTTPException(status_code=404, detail="结果文件不存在")
    
    return FileResponse(
        result_path,
        media_type="video/mp4",
        filename=os.path.basename(result_path)
    )

@app.get("/api/models")
async def get_models_info():
    """获取模型信息"""
    return {
        "text2video": {
            "available": generator is not None,
            "model": "Stable Diffusion v1.5" if generator else None
        },
        "style_transfer": {
            "available": transfer is not None,
            "model": "StyleGAN3" if transfer else None
        },
        "video_editing": {
            "available": editor is not None,
            "model": "Video-LLaMA" if editor else None
        }
    }

@app.post("/api/load_models")
async def load_models_api():
    """加载模型API"""
    success = load_models()
    return {
        "success": success,
        "message": "模型加载成功" if success else "模型加载失败"
    }

@app.get("/api/supported_styles")
async def get_supported_styles():
    """获取支持的风格列表"""
    return {
        "text2video": [
            "realistic", "cartoon", "anime", "oil_painting",
            "watercolor", "sketch", "cyberpunk", "vintage"
        ],
        "style_transfer": [
            "cartoon", "anime", "oil_painting", "watercolor",
            "sketch", "cyberpunk", "vintage"
        ]
    }

@app.get("/api/supported_operations")
async def get_supported_operations():
    """获取支持的编辑操作"""
    return {
        "operations": [
            "trim", "resize", "crop", "speed", "fade",
            "blur", "brightness", "contrast", "saturation",
            "noise", "text", "logo", "transition", "stabilize", "denoise"
        ]
    }

def create_fastapi_app():
    """创建FastAPI应用"""
    return app

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000) 